package utils.global.limitflow;

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import utils.global.constant.RestfulJson;
import utils.global.constant.RestfulType;
import utils.global.utils.TimeUtils;

import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;

@Aspect
@Slf4j
public class LimitFlowAop {
    private ConcurrentHashMap<String, RateLimiter> POOL = new ConcurrentHashMap<>();

    @Pointcut("@annotation(limitFLow)")
    public void Limit(LimitFlow limitFLow) {
    }

    @Around("Limit(limitFlow)")
    public Object LimitAspect(ProceedingJoinPoint point, LimitFlow limitFlow) throws Throwable {

        Method method = ((MethodSignature) point.getSignature()).getMethod();

        LimitFlow limitFLow = method.getAnnotation(LimitFlow.class);

        String name = limitFLow.name();

        if (!POOL.containsKey(name)) {
            POOL.put(name, RateLimiter.create(limitFLow.maxFlow(), limitFLow.expandedTime(), limitFLow.timeUnit()));
        } else {
            if (!POOL.get(name).tryAcquire()) {
                log.warn(TimeUtils.sdf.format(System.currentTimeMillis()) + " : 流量策略" + name + " 限流中");
                return RestfulJson.type(RestfulType.FLOW_AILURE);
            }
        }

        Object proceed = point.proceed();

        return proceed;
    }
}
